import java.io.File;
import java.io.FileNotFoundException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Scanner;
import java.util.Set;
import java.util.TreeSet;



public class PosTagger 
{
	private static final String begin = "<S>";
	private static final String end = "<E>";
	private static final double filesize = 760341.0;
	private static int countLines= 0; 
	private static double precisionCount =0.0;
	private static double countAllOutputted = 0.0;
	private static HashMap<Trigram, Double> trigramToPrbs= new HashMap<Trigram, Double>();
	private static HashMap<Bigram, Double> bigramToPrbs = new HashMap<Bigram, Double>();
	private static HashMap<Unigram, Double> unigramToPrbs = new HashMap<Unigram, Double>();
	private static HashMap<Trigram, Double> trigramToSmoothedPrbs= new HashMap<Trigram, Double>();
	private static HashMap<Bigram, Double> bigramToSmoothedPrbs = new HashMap<Bigram, Double>();
	private static HashMap<Bigram, Double> bigramAlphas = new HashMap<Bigram, Double>();
	private static HashMap<Unigram, Double> unigramAlphas = new HashMap<Unigram, Double>();
	private static HashMap<String, Double> posToFreq = new HashMap<String, Double>();
	private static HashMap<String, Set<String>> wordToAllPOS = new HashMap<String, Set<String>>();
	private static ArrayList<String> testSentences = new ArrayList<String>();
	private static ArrayList<ArrayList<String>> testSentencesList = new ArrayList<ArrayList<String>>();
	private static ArrayList<String> correctTags = new ArrayList<String>();
	private static ArrayList<ArrayList<String>> correctTagsList = new ArrayList<ArrayList<String>>();
	private static ArrayList<ArrayList<ViterbiProbs>>  listOfViterbiProbsObjects;
	private static HashMap<WordAndTag, Double> wordAndTagToUnsmoothedPrb = 
		new HashMap<WordAndTag, Double>();
	public static void main(String[] args)
	{
		//a map for wordAndTag to frequency
		HashMap<WordAndTag, Double> wordAndTagToFreq = new HashMap<WordAndTag, Double>();
		//map sfor ngram to frequency
		HashMap<Trigram, Double> trigramToFreq = new HashMap<Trigram, Double>();
		HashMap<Bigram, Double> bigramToFreq = new HashMap<Bigram, Double>();
		HashMap<Unigram, Double> unigramToFreq = new HashMap<Unigram, Double>();
		System.out.println("The program is reading the file...");
		//extracting the information relevant later for probability estimations
		readFile(wordAndTagToFreq, trigramToFreq,bigramToFreq,unigramToFreq, posToFreq);
		System.out.println("The program is calculating lexical and POS probabilities...");
		calculateUnsmoothedLexPrbs(wordAndTagToFreq,wordAndTagToUnsmoothedPrb, posToFreq);
		calculateUnsmoothedPosPrbs(unigramToFreq, bigramToFreq, trigramToFreq);
		calculateAlphas();
		System.out.println("The program is calculating the best POS sequence...this might take a while");
		annotate();
	}

	//reads the file and extracts information relevant later for probability estimations
	private static void readFile(HashMap<WordAndTag, Double> wordAndTagToFreq, 
			HashMap<Trigram, Double> trigramToFreq,
			HashMap<Bigram, Double> bigramToFreq,
			HashMap<Unigram, Double> unigramToFreq,
			HashMap<String, Double> posToFreq)
	{
		//String fileName="data/tiger";
		String fileName="tiger";
		Scanner input = null;
		try
		{
			input = new Scanner (new File(fileName));
		}
		catch(FileNotFoundException e)
		{
			System.out.println("Could not open the file " + fileName);
			System.exit(0);
		}
		while((input.hasNextLine()) && (countLines<=(0.9*filesize)))
		{
			//beggining of a sentence
			String line = input.nextLine();
			countLines++;
			String[] tmpArr = line.split("\\t");
			String pos1 = begin;
			String pos2 = begin;
			String pos3 = tmpArr[1];
			updatePOSMap(posToFreq, tmpArr[1]);
			//creaing a wordAndTag object and storing in a map
			WordAndTag wordAndTag = WordAndTag.newWordAndTag(tmpArr[0], tmpArr[1]);
			updateWordAndTagMap(wordAndTagToFreq, wordAndTag);
			//creating ngram objects and storing in maps

			Bigram bigram = Bigram.newBigram(pos1, pos2);
			Unigram unigram = Unigram.newUnigram(pos1);
			updateBigramMap(bigramToFreq, bigram);
			countLines += 2.0;
			updateUnigramMap(unigramToFreq, unigram);
			updateUnigramMap(unigramToFreq, unigram);
			Trigram trigram = Trigram.newTrigram(pos1, pos2, pos3);
			bigram = Bigram.newBigram(pos2, pos3);
			unigram = Unigram.newUnigram(pos3);
			//System.out.println("start of sentence: " +trigram);
			updateTrigramMap(trigramToFreq, trigram);
			updateBigramMap(bigramToFreq, bigram);
			updateUnigramMap(unigramToFreq, unigram);
			//sentence 
			while(!line.equals(""))
			{
				line = input.nextLine();
				countLines++;
				if(line.equals(""))
					break;
				//System.out.println(line);
				tmpArr = line.split("\\t");
				pos1 = pos2;
				pos2 = pos3;
				pos3 = tmpArr[1];
				updatePOSMap(posToFreq, tmpArr[1]);
				wordAndTag = WordAndTag.newWordAndTag(tmpArr[0], tmpArr[1]);
				updateWordAndTagMap(wordAndTagToFreq, wordAndTag);
				trigram = Trigram.newTrigram(pos1, pos2, pos3);
				bigram = Bigram.newBigram(pos2, pos3);
				unigram = Unigram.newUnigram(pos3);
				//System.out.println("middle of sentence: " +trigram);
				updateTrigramMap(trigramToFreq, trigram);
				updateBigramMap(bigramToFreq, bigram);
				updateUnigramMap(unigramToFreq, unigram);
			}
			//end of a sentence
			pos1 = pos2;
			pos2 = pos3;
			pos3 = end;
			updatePOSMap(posToFreq, pos3);
			wordAndTag = WordAndTag.newWordAndTag(pos3, pos3);
			updateWordAndTagMap(wordAndTagToFreq, wordAndTag);
			trigram = Trigram.newTrigram(pos1, pos2, pos3);
			bigram = Bigram.newBigram(pos2, pos3);
			//unigram = Unigram.newUnigram(pos1);
			unigram = Unigram.newUnigram(pos3);
			//System.out.println("end of sentence: " +trigram);
			updateTrigramMap(trigramToFreq, trigram);
			updateBigramMap(bigramToFreq, bigram);
			updateUnigramMap(unigramToFreq, unigram);
		}
		System.out.println("size of read file:" + countLines);
		//store remaining text in an array
		while (input.hasNextLine())
		{
			while(true)
			{
				String line = input.nextLine();
				if(line.equals(""))
					break;
				else
				{
					String[] tmpArr = line.split("\\t");
					testSentences.add(tmpArr[0]);
					correctTags.add(tmpArr[1]);
				}
			}
			//System.out.println(testSentences);
			ArrayList<String> tmpCopy = (ArrayList<String>) testSentences.clone();
			testSentencesList.add(tmpCopy);
			testSentences.clear();
			ArrayList<String> tmpCopyCorrectTags = (ArrayList<String>) correctTags.clone();
			correctTagsList.add(tmpCopyCorrectTags);
			correctTags.clear();
		}

	}

	//updates the map containing trigrams as keys and their respective frequencies as values
	private static void updateTrigramMap(HashMap<Trigram, Double> trigramToFreq,
			Trigram trigram)
	{
		if(!trigramToFreq.containsKey(trigram))
		{
			trigramToFreq.put(trigram, 1.0);
		}
		else
		{
			//System.out.println("Found in map!");
			Double value = trigramToFreq.get(trigram)+ 1.0;
			trigramToFreq.put(trigram, value);
		}
	}
	//updates the map containing bigrams as keys and their respective frequencies as values
	private static void updateBigramMap(HashMap<Bigram, Double> bigramToFreq,
			Bigram bigram)
	{
		if(!bigramToFreq.containsKey(bigram))
		{
			bigramToFreq.put(bigram, 1.0);
		}
		else
		{
			//System.out.println("Found in map!");
			Double value = bigramToFreq.get(bigram)+ 1.0;
			bigramToFreq.put(bigram, value);
		}
	}
	//updates the map containing bigrams as keys and their respective frequencies as values
	private static void updateUnigramMap(HashMap<Unigram, Double> unigramToFreq,
			Unigram unigram)
	{
		if(!unigramToFreq.containsKey(unigram))
		{
			unigramToFreq.put(unigram, 1.0);
		}
		else
		{
			//System.out.println("Found in map!");
			Double value = unigramToFreq.get(unigram)+ 1.0;
			unigramToFreq.put(unigram, value);
		}
	}
	//update the map containing pos as keys and their frequencies as values
	private static void updatePOSMap(HashMap<String, Double> posToFreq,
			String pos)
	{
		if(!posToFreq.containsKey(pos))
		{
			posToFreq.put(pos, 1.0);
		}
		else
		{
			Double value = posToFreq.get(pos)+ 1.0;
			posToFreq.put(pos, value);
		}

	}

	//update the map containing wordAndTags as keys and their frequencies as values
	private static void updateWordAndTagMap(HashMap<WordAndTag, Double> wordAndTagToFreq,
			WordAndTag wordAndTag)
	{

		String keyForWordToPosMap= wordAndTag.getWord();
		String valueToAddToPosMap= wordAndTag.getPOS();
		Set<String> setToFill=null;
		if(wordToAllPOS.get(keyForWordToPosMap)!=null)
		{
			setToFill= wordToAllPOS.get(keyForWordToPosMap);
		}
		else
		{
			setToFill = new TreeSet<String>();
		}
		setToFill.add(valueToAddToPosMap);
		wordToAllPOS.put(keyForWordToPosMap, setToFill);

		if(!wordAndTagToFreq.containsKey(wordAndTag))
		{
			wordAndTagToFreq.put(wordAndTag, 1.0);
		}
		else
		{
			Double value = wordAndTagToFreq.get(wordAndTag)+ 1.0;
			wordAndTagToFreq.put(wordAndTag, value);
		} 
	}
	//calculating unsmooted lexical probabilities
	private static void calculateUnsmoothedLexPrbs(HashMap<WordAndTag, Double> wordAndTagToFreq,
			HashMap<WordAndTag, Double> wordAndTagToUnsmoothedPrb, 
			HashMap<String, Double> posToFreq)
	{
		for(WordAndTag wordAndTag: wordAndTagToFreq.keySet())
		{
			double value = wordAndTagToFreq.get(wordAndTag) / posToFreq.get(wordAndTag.getPOS());
			wordAndTagToUnsmoothedPrb.put(wordAndTag, value);
		}

	}


	//calculating smoothed pos probabilities for trigrams
	private static void calculateSmoothedPosPrbsTrigram(Trigram trigram)
	{
		if(!trigramToSmoothedPrbs.containsKey(trigram))
		{
			Bigram bigramAlpha = Bigram.newBigram(trigram.getPOS1(), trigram.getPOS2());
			Bigram bigramPrb = Bigram.newBigram(trigram.getPOS2(), trigram.getPOS3());
			double bigramAlphaPrb = 1.0;
			double bigramPrbSmoothed=1.0;
			//System.out.println(bigramAlphas.get(bigramAlpha) + " " + bigramToPrbs.get(bigramPrb));
			if(bigramAlphas.get(bigramAlpha)!=null)
				bigramAlphaPrb = bigramAlphas.get(bigramAlpha);
			if(bigramToPrbs.get(bigramPrb)!=null)
				bigramPrbSmoothed= bigramToPrbs.get(bigramPrb);
			else 
			{
				calculateSmoothedPosPrbsBigram(bigramPrb);
				bigramPrbSmoothed= bigramToSmoothedPrbs.get(bigramPrb);
			}
			double valueSmoothed = bigramAlphaPrb * bigramPrbSmoothed;
			trigramToSmoothedPrbs.put(trigram, valueSmoothed);
		}


	}

	//calculating smoothed pos probabilities for bigrams
	private static void calculateSmoothedPosPrbsBigram(Bigram bigram)
	{
		if(!bigramToSmoothedPrbs.containsKey(bigram))
		{
			Unigram unigramAlpha = Unigram.newUnigram(bigram.getPOS1());
			Unigram unigramPrb = Unigram.newUnigram(bigram.getPOS2());
			double valueSmoothed = unigramAlphas.get(unigramAlpha) * unigramToPrbs.get(unigramPrb);
			bigramToSmoothedPrbs.put(bigram, valueSmoothed);
		}


	}


	//calculating alphass
	private static void calculateAlphas()
	{
		double sigmaA= 0.0;
		double sigmaB= 0.0;
		//calculating unigram alphas
		for(Unigram unigram: unigramToPrbs.keySet())
		{
			String alphaString = unigram.getPOS1();
			for(Bigram bigram: bigramToPrbs.keySet())
			{
				if(bigram.getPOS1().equals(alphaString)) {
					sigmaA += bigramToPrbs.get(bigram);
					//if(unigram.getPOS1().equals(bigram.getPOS2()))
					Unigram unigram2 = Unigram.newUnigram(bigram.getPOS2()) ;
					sigmaB += unigramToPrbs.get(unigram2);
				}
			}
			if(sigmaA>1 || sigmaB>1)
				System.out.println("sigma A unigram: " + sigmaA + " unigram context: " + 
						unigram +   " sigma B unigram: " + sigmaB);   
			Double value = (1.0 - sigmaA) / (1.0 - sigmaB); 
			unigramAlphas.put(unigram, value);
			sigmaA=0.0;
			sigmaB=0.0;
		}

		//calculating bigram alphas
		for(Bigram bigram: bigramToPrbs.keySet())
		{
			sigmaA= 0.0;
			sigmaB= 0.0;
			String pos1 = bigram.getPOS1();
			String pos2 = bigram.getPOS2();
			for(Trigram trigram: trigramToPrbs.keySet())
			{
				if( (trigram.getPOS1().equals(pos1)) &&  (trigram.getPOS2().equals(pos2)) )
					sigmaA += trigramToPrbs.get(trigram);
			}

			for(Bigram bigram2: bigramToPrbs.keySet())
			{

				if( (bigram2.getPOS1().equals(pos2)))
				{
					sigmaB += bigramToPrbs.get(bigram2);    

				}
			}

			Double value = (1.0 - sigmaA) / (1.0 - sigmaB); 
			if(sigmaA>1 || sigmaB>1)
				System.out.println("Bigram: " + bigram + " sigma A: " + sigmaA + " sigma B: " + sigmaB);
			bigramAlphas.put(bigram, value);
			sigmaA=0.0;
			sigmaB=0.0;

		}

	}

	//calculating unsmooted pos probabilities
	private static void calculateUnsmoothedPosPrbs(HashMap<Unigram, Double> unigramToFreq, 
			HashMap<Bigram, Double>bigramToFreq, 
			HashMap<Trigram, Double> trigramToFreq)
	{
		double sum= 0.0;
		for(Trigram trigram: trigramToFreq.keySet())
		{
			String[] tmp =trigram.getTrigram();
			Bigram bigram = Bigram.newBigram(tmp[0], tmp[1]);
			Double value = bigramToFreq.get(bigram);
			//	    System.out.println("trigramFreq"+ trigramToFreq.get(trigram) 
			//		    + " value: " + value+ " bigram: " + bigram);
			Double valueToPut = (trigramToFreq.get(trigram) - 0.5) / value;
			trigramToPrbs.put(trigram, valueToPut);
			sum += valueToPut;
		}
		//System.out.println("sum: " + sum);
		sum= 0.0;
		for(Bigram bigram: bigramToFreq.keySet())
		{
			String[] tmp =bigram.getBigram();

			Unigram unigram = Unigram.newUnigram(tmp[0]);

			Double value = unigramToFreq.get(unigram);
			Double valueToPut = (bigramToFreq.get(bigram) - 0.5) / value;
			//  if(unigram.getPOS1().equals("VMINF"))
			//		System.out.println("prb is: " + valueToPut + " freq:" + bigramToFreq.get(bigram)
			//			+" unigram freq: " + value
			//			+ " bigram: " + bigram.toString());
			bigramToPrbs.put(bigram, valueToPut);

			sum += valueToPut;
		}
		//	System.out.println("sum: " + sum);

		sum=0.0;
		for(Unigram unigram: unigramToFreq.keySet())
		{
			Double value = unigramToFreq.get(unigram);
			Double valueToPut = value / countLines;
			unigramToPrbs.put(unigram, valueToPut);
			sum += valueToPut;
			//System.out.println(bigram + " "+ value);
		}
		//	System.out.println("sum: " + sum);
	}


	//calculating viterbi probabilities
	private static void annotate()
	{
		int count =0;
		System.out.println("How many sentences would you like to be printed out?\n" +
				"Total number of sentences in the training is: " + testSentencesList.size());
		Scanner keyboard = new Scanner(System.in);
		int numSentencesToOutput = keyboard.nextInt();
		double sumOfFrequenciesForPOS = 0.0;
		for(String posTag: posToFreq.keySet())
		{
			sumOfFrequenciesForPOS += posToFreq.get(posTag);
		}
		for(int j=0; j<numSentencesToOutput; j++)
		{
			System.out.println("Sentence number: " + (j+1));
			testSentences = testSentencesList.get(j);
			correctTags = correctTagsList.get(j);
			//System.out.println(testSentences);
			//v[i,t,t'] := 0 for all tags t, t' and 0 <= i <= n + 1
			listOfViterbiProbsObjects= new ArrayList<ArrayList<ViterbiProbs>>(testSentences.size()+1);
			for(int i=0; i<testSentences.size()+1; i++)
			{
				listOfViterbiProbsObjects.add(new ArrayList<ViterbiProbs>());
			}
			//v[0,<S>, <S>] := 1;
			ViterbiProbs start= ViterbiProbs.newViterbiProb(begin, begin, 1.0, "");
			ArrayList<ViterbiProbs> viterbiProbsTmpList = new ArrayList<ViterbiProbs>();
			viterbiProbsTmpList.add(start);
			listOfViterbiProbsObjects.set(0, viterbiProbsTmpList);
			//For i := 1 to n
			for(int i=1; i<testSentences.size(); i++)
			{
				String word= testSentences.get(i);
				//System.out.println(word);
				ArrayList<ViterbiProbs> tmp = new ArrayList<ViterbiProbs>();

				if(wordToAllPOS.get(word)!=null)
				{
					for(String posOfWord: wordToAllPOS.get(word))
					{
						//process(i,p(wi|t),t)
						WordAndTag wordAndTag =WordAndTag.newWordAndTag(word, posOfWord);
						process(i, wordAndTagToUnsmoothedPrb.get(wordAndTag), wordAndTag.getPOS(),tmp);
					}
				}
				else
				{
					for(String pos: posToFreq.keySet())
					{
						//restricting to content words
						if((pos.startsWith("NN")) || (pos.startsWith("JJ")) || (pos.startsWith("VB")))
							process(i, (posToFreq.get(pos)/sumOfFrequenciesForPOS), pos,tmp);
					}
				}
				//System.out.println(i +": "+listOfViterbiProbsObjects.get(i)); 
			}
			//process(n+1,1,<E>)
			ArrayList<ViterbiProbs> tmp = new ArrayList<ViterbiProbs>();
			process(testSentences.size(), 1.0, end, tmp);
			//System.out.println(testSentences.size() +": "+listOfViterbiProbsObjects.get(testSentences.size())); 
			//t := arg maxt v[n + 1; t; <E>]
			String posMaxT = "";
			double maxPrb=0.0;	
			//	System.out.println("listOfViterbiProbsObjects.get(listOfViterbiProbsObjects.size()-1): " 
			//		+ listOfViterbiProbsObjects.get(listOfViterbiProbsObjects.size()-1));
			for(ViterbiProbs viterbiPrb: listOfViterbiProbsObjects.get(testSentences.size()))
			{
				if(viterbiPrb.getProb()>maxPrb)
				{
					maxPrb= viterbiPrb.getProb();
					posMaxT= viterbiPrb.getBestPreviousTag();
				}
			}
			//print(n,t,<E>)
			print(testSentences.size()-1, posMaxT, "$.");
			System.out.println("_________________________________________________________________");
		}

		System.out.println("Precision is: " + ((precisionCount/countAllOutputted)*100)
				+ "%");
	}


	private static void process(int index, double prb, String pos, ArrayList<ViterbiProbs> tmp)
	{
		int vCurrentColumnIndex=-1;
		ViterbiProbs vprobCurrent = null;
		for(ViterbiProbs vprob: listOfViterbiProbsObjects.get(index-1))
		{
			Trigram trigram = Trigram.newTrigram(vprob.getPOS1(), vprob.getPOS2(), pos);
			double v=0.0;
			if(trigramToPrbs.get(trigram)!=null)
			{
				v = vprob.getProb()  * prb * trigramToPrbs.get(trigram);
			}
			else
			{
				calculateSmoothedPosPrbsTrigram(trigram);
				v = vprob.getProb() * prb * trigramToSmoothedPrbs.get(trigram);
			}
			for(int i=0; i< listOfViterbiProbsObjects.get(index).size(); i++)
			{
				if( (listOfViterbiProbsObjects.get(index).get(i).getPOS1().equals(vprob.getPOS2())) 
						&& (listOfViterbiProbsObjects.get(index).get(i).getPOS2().equals(pos)))
				{  
					vCurrentColumnIndex = i;
				}
			}
			if (vCurrentColumnIndex!=-1)
			{
				double valueToPut = 0.0;
				if (v > listOfViterbiProbsObjects.get(index).get(vCurrentColumnIndex).getProb())
				{
					valueToPut = v;
				}
				else
					valueToPut = listOfViterbiProbsObjects.get(index).get(vCurrentColumnIndex).getProb(); 
				vprobCurrent= 
					ViterbiProbs.newViterbiProb(vprob.getPOS2(), pos,  valueToPut, vprob.getPOS1());
				listOfViterbiProbsObjects.get(index).set(vCurrentColumnIndex, vprobCurrent);

			}
			else
			{
				ViterbiProbs vCurrent= 
					ViterbiProbs.newViterbiProb(vprob.getPOS2(), pos,  v, vprob.getPOS1());
				listOfViterbiProbsObjects.get(index).add(vCurrent);
			}
		}
	}

	//print(i,t,t')
	private static void print(int i, String posT, String posTprime)
	{
		String bestPrevious= posT;
		String pos2= posTprime;
		//if i > 1 
		if(i>1)
		{
			//a[i+1,t,t']
			for(ViterbiProbs viterbiPrb: listOfViterbiProbsObjects.get(i))
			{
				if( (viterbiPrb.getPOS1().equals(posT)) && 
						(viterbiPrb.getPOS2().equals(posTprime)) )
				{
					bestPrevious = viterbiPrb.getBestPreviousTag();
					pos2 = viterbiPrb.getPOS1();
				}
			}
			//print(i-1,a[i+1,t,t'],t)
			print(i-1, bestPrevious, pos2);
			//precision count
			if(posTprime.equals(correctTags.get(i)))
				precisionCount+=1.0;
			countAllOutputted+=1.0;
			//output wi; t	
			System.out.println(testSentences.get(i).toString() + " [" + posTprime + "]"
					+ "  correct: {" + correctTags.get(i) + "}");
		}
	}

}
